import torch.nn as nn
from Fast_RCNN_Extractor import get_device
from collections import OrderedDict


class Fully_Conection_Model(nn.Module):
    """
    全连接网络 , 其中一个做classification 一个做线性回归
    """
    def __init__(self):
        super(Fully_Conection_Model, self).__init__()
        self.First_Sibling_Layer = nn.Sequential(
            OrderedDict([('Dropout_1', nn.Dropout()),
                         ('Linear_1', nn.Linear(512*7*7, 4096)),
                         ('ReLU_1', nn.ReLU(inplace=True)),

                         ('Dropout_2', nn.Dropout()),
                         ('Linear_2', nn.Linear(4096, 4096)),
                         ('ReLU_2', nn.ReLU(inplace=True)),

                         ('Linear_3', nn.Linear(4096, 21)),
                         ])
        )

        self.Second_Sibling_Layer = nn.Sequential(
            OrderedDict([('Dropout_1', nn.Dropout()),
                         ('Linear_1', nn.Linear(512 * 7 * 7, 4096)),
                         ('ReLU_1', nn.ReLU(inplace=True)),

                         ('Dropout_2', nn.Dropout()),
                         ('Linear_2', nn.Linear(4096, 4096)),
                         ('ReLU_2', nn.ReLU(inplace=True)),

                         ('Linear_3', nn.Linear(4096, 4)),
                         ])
        )

    def forward(self, input):
        input = input.view(-1, 512*7*7)
        class_out = self.First_Sibling_Layer(input)
        box_out = self.Second_Sibling_Layer(input)
        return class_out, box_out


Fast_RCNN_Model = Fully_Conection_Model().to(device=get_device())